import torch 
from typing import Optional 
from pydantic import BaseModel 
from fastapi import FastAPI, HTTPException, Header 
import uvicorn 
from transformers import (
    AutoModelForCausalLM, AutoTokenizer, AutoConfig, GenerationConfig, TextIteratorStreamer 
) 
from fastapi.responses import StreamingResponse 
import threading
import os

# 加载分词器
def load_math_tokenizer(model_name_or_path, fast_tokenizer=True): 
    tokenizer = AutoTokenizer.from_pretrained(
        model_name_or_path, fast_tokenizer=fast_tokenizer, trust_remote_code=True) 
    # 数学模型通常需要明确的特殊符号处理
    tokenizer.pad_token = tokenizer.eos_token if tokenizer.pad_token is None else tokenizer.pad_token
    tokenizer.bos_token = tokenizer.eos_token if tokenizer.bos_token is None else tokenizer.bos_token
    tokenizer.padding_side = 'left'  # 左侧填充更适合生成任务
    tokenizer.truncation_side = "left"  # 超长时截断左侧
    return tokenizer 

# 加载数学解题模型
def create_math_model(model_name_or_path): 
    model_config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
    model_config.use_cache = True
    model = AutoModelForCausalLM.from_pretrained(
        model_name_or_path, 
        config=model_config, 
        trust_remote_code=True,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    ) 
    return model 

# 定义请求数据模型
class MathSolveRequest(BaseModel): 
    problem: str  # 数学问题描述
    max_input_length: Optional[int] = 1024  # 数学问题通常不需要太长的输入长度
    max_new_tokens: Optional[int] = 8192  # 解题步骤可能需要较长输出
    temperature: Optional[float] = 0.6  # 数学推理适当提高随机性
    top_p: Optional[float] = 0.95  # 控制生成多样性
    do_sample: Optional[bool] = True 

# 初始化FastAPI应用
app = FastAPI(title="Math Solver API") 
model = None 
tokenizer = None 
device = None 

# 从环境变量获取API密钥
API_TOKEN = os.environ.get("MATH_API_TOKEN", None) 

# 启动时加载模型
@app.on_event("startup") 
def startup_event(): 
    global model, tokenizer, device, gen_base_cfg 
    # 数学模型默认路径，可通过环境变量覆盖
    base_model = os.environ.get("MATH_MODEL_PATH", "/data/FAITA/yzc/verl_hf/step480") 
    print(f"Loading math solver model: {base_model}") 
    
    # 配置GPU
    gpu_id = os.environ.get("MATH_GPU_ID", "5") 
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id 
    
    # 加载分词器和模型
    tokenizer = load_math_tokenizer(base_model) 
    model = create_math_model(base_model) 
    device = torch.device(f"cuda:0" if torch.cuda.is_available() else "cpu") 
    
    # 模型精度设置
    try:
        model = model.half().to(device)  # 半精度加速推理
    except Exception as e:
        print(f"Failed to use half precision: {e}, using float32 instead")
        model = model.to(device)
    
    model.eval()  # 推理模式
    
    # 数学推理基础配置
    gen_base_cfg = { 
        "do_sample": True, 
        "top_p": 0.95, 
        "temperature": 0.6, 
        "num_return_sequences": 1 
    } 
    print(f"Math model ready. Device: {device}, Using GPU ID: {gpu_id}") 

# 将数据移至指定设备
def to_device(batch, device): 
    output = {} 
    for k, v in batch.items(): 
        try:
            output[k] = v.to(device)
        except:
            output[k] = v
    return output 

# 非流式解题接口
@app.post("/solve") 
def solve_math(req: MathSolveRequest, authorization: Optional[str] = Header(None)): 
    # 认证检查
    if API_TOKEN: 
        if not authorization or authorization != f"Bearer {API_TOKEN}": 
            raise HTTPException(status_code=401, detail="Unauthorized: Invalid or missing token") 

    # 构建数学问题提示（可根据模型要求调整格式）
    messages = [
        {
         "role": "system",  # 必须是 system 角色，优先级最高
        "content": """你是严格遵循格式要求的数学解题专家，以下规则**必须无条件遵守**，否则回答无效：

    1. **语言强制**：所有内容（包括思考过程、步骤、结论）只能用中文，禁止出现任何英文（公式符号除外）。
    - 错误示例："To solve this equation, we need to..."（禁止）
    - 正确示例："为求解该方程，我们需要..."（必须）

    2. **公式格式强制**：所有数学公式、符号、表达式必须用 $ 包裹：
    - 行内公式：用单个 $ 包裹（如：速度 $v = s/t$ m/s）
    - 块级公式：用两个 $ 包裹（如：$$\\int_0^1 x^2 dx = \\frac{1}{3}$$）
    - 错误示例：$\\int_0^1 x^2 dx$（块级公式用了单 $，错误）、v = s/t（未用 $ 包裹，错误）
    - 正确示例：$$\\int_0^1 x^2 dx$$（块级用双 $）、$v = s/t$（行内用单 $）
    """
        },
        {    "role": "user", 
            "content": f"""你是一个专业的数学解题助手，需遵循以下规则：
                1. 所有思考过程、推理步骤和最终结论必须使用**中文**表达。
                2. 涉及的数学公式（包括符号、表达式、方程等）必须用`$`包裹：
                - 行内公式用单个`$`包裹（如：$a + b = c$);
                - 块级公式用两个`$`包裹（如：$$\\int_0^1 x^2 dx$$）。
                3. 解题时先生成简洁的推理步骤链，格式为<chain>步骤1 -> 步骤2 -> ... -> 最终步骤</chain>，每个步骤需简短且概括核心操作。
                4. 按照步骤链逐步展开解答，每一步需验证正确性，若有疑问可重新生成步骤链并重试。
{req.problem}"""
        }
    ]
    # 自动生成符合DeepSeek格式的prompt
    prompt = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True  # 自动添加"assistant: "前缀
    )

    # 处理输入
    inputs = tokenizer( 
        prompt, 
        return_tensors="pt", 
        padding=True, 
        truncation=True, 
        max_length=req.max_input_length 
    ) 
    inputs = to_device(inputs, device) 

    # 生成配置
    gen_cfg = GenerationConfig( 
        temperature=req.temperature, 
        top_p=req.top_p, 
        do_sample=req.do_sample, 
        num_return_sequences=1 
    ) 
    
    # 模型推理
    with torch.no_grad(): 
        outs = model.generate( 
            **inputs, 
            generation_config=gen_cfg, 
            max_new_tokens=req.max_new_tokens, 
            pad_token_id=tokenizer.eos_token_id 
        ) 
    
    # 解码结果并去除提示部分
    full_text = tokenizer.decode(outs[0], skip_special_tokens=True)
    solution = full_text.replace(prompt, "").strip()
    return {"problem": req.problem, "solution": solution} 

# 流式解题接口
@app.post("/solve_stream") 
def solve_math_stream(req: MathSolveRequest, authorization: Optional[str] = Header(None)): 
    # 认证检查
    if API_TOKEN: 
        if not authorization or authorization != f"Bearer {API_TOKEN}": 
            raise HTTPException(status_code=401, detail="Unauthorized: Invalid or missing token") 

    # 构建提示
    messages = [
        {
         "role": "system",  # 必须是 system 角色，优先级最高
        "content": """你是严格遵循格式要求的数学解题专家，以下规则**必须无条件遵守**，否则回答无效：

    1. **语言强制**：所有内容（包括思考过程、步骤、结论）只能用中文，禁止出现任何英文（公式符号除外）。
    - 错误示例："To solve this equation, we need to..."（禁止）
    - 正确示例："为求解该方程，我们需要..."（必须）
    """
        },
        {    "role": "user", 
            "content": f"""你是一个专业的数学解题助手，需遵循以下规则：
                1. 所有思考过程、推理步骤和最终结论必须使用**中文**表达。
                3. 解题时先生成简洁的推理步骤链，格式为<chain>步骤1 -> 步骤2 -> ... -> 最终步骤</chain>，每个步骤需简短且概括核心操作。
                4. 按照步骤链逐步展开解答，每一步需验证正确性，若有疑问可重新生成步骤链并重试。
{req.problem}"""
        }
    ]
    # 自动生成符合DeepSeek格式的prompt
    prompt = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True  # 自动添加"assistant: "前缀
    )
    
    # 处理输入
    inputs = tokenizer( 
        prompt, 
        return_tensors="pt", 
        padding=True, 
        truncation=True, 
        max_length=req.max_input_length 
    ) 
    inputs = to_device(inputs, device) 

    # 流式输出配置
    streamer = TextIteratorStreamer(
        tokenizer, 
        skip_special_tokens=True, 
        timeout=15.0,  # 数学推理可能需要更长时间
        skip_prompt=True  # 跳过提示部分的输出
    ) 

    # 启动后台生成线程
    threading.Thread( 
        target=model.generate, 
        kwargs={ 
            **inputs, 
            "max_new_tokens": req.max_new_tokens, 
            "temperature": req.temperature, 
            "top_p": req.top_p, 
            "do_sample": req.do_sample, 
            "streamer": streamer, 
            "pad_token_id": tokenizer.eos_token_id 
        }, 
        daemon=True 
    ).start() 

    # 返回流式响应
    return StreamingResponse(streamer, media_type="text/plain") 

# 启动服务
if __name__ == "__main__": 
    # 数学服务使用8011端口，避免与其他模型冲突
    uvicorn.run("math_deploy:app", host="10.143.12.79", port=8077, workers=1)